{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "ctaLvRUsfpj8",
      "metadata": {
        "id": "ctaLvRUsfpj8"
      },
      "source": [
        "<a target=\"_blank\" href=\"https://colab.research.google.com/github/cohere-ai/notebooks/blob/main/notebooks/llmu/RAG_with_Chat_Embed_and_Rerank.ipynb\">\n",
        "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
        "</a>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "61bac5b5",
      "metadata": {
        "id": "61bac5b5"
      },
      "source": [
        "# RAG with Chat, Embed, and Rerank\n",
        "\n",
        "This notebook shows how to build a RAG-powered chatbot with Cohere's Chat endpoint.  The chatbot can extract relevant information from external documents and produce verifiable, inline citations in its responses.\n",
        "\n",
        "This application will use several Cohere API endpoints:\n",
        "\n",
        "- Chat: For handling the main logic of the chatbot, including turning a user message into queries, generating responses, and producing citations\n",
        "- Embed: For turning textual documents into their embeddings representation, later to be used in retrieval (we’ll use the latest, state-of-the-art Embed v3 model)\n",
        "- Rerank: For reranking the retrieved documents according to their relevance to a query\n",
        "\n",
        "The diagram below provides an overview of what we’ll build."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "CLNMDV4IuA3C",
      "metadata": {
        "id": "CLNMDV4IuA3C"
      },
      "source": [
        "![rag-workflow-2.png]()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f6ab2d5d",
      "metadata": {
        "id": "f6ab2d5d"
      },
      "source": [
        "Here is a summary of the steps involved.\n",
        "\n",
        "Initial phase:\n",
        "- **Step 0**: Ingest the documents – get documents, chunk, embed, and index.\n",
        "\n",
        "For each user-chatbot interaction:\n",
        "- **Step 1**: Get the user message\n",
        "- **Step 2**: Call the Chat endpoint in query-generation mode\n",
        "- If at least one query is generated\n",
        "  - **Step 3**: Retrieve and rerank relevant documents\n",
        "  - **Step 4**: Call the Chat endpoint in document mode to generate a grounded response with citations\n",
        "- If no query is generated\n",
        "  - **Step 4**: Call the Chat endpoint in normal mode to generate a response"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "TWyo_5WoNUM-",
      "metadata": {
        "id": "TWyo_5WoNUM-"
      },
      "source": [
        "# Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5pLAhQmTOKiV",
      "metadata": {
        "id": "5pLAhQmTOKiV"
      },
      "outputs": [],
      "source": [
        "! pip install cohere hnswlib unstructured -q"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "id": "f3a03a57",
      "metadata": {
        "id": "f3a03a57"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "import cohere\n",
        "from pinecone import Pinecone, PodSpec\n",
        "import uuid\n",
        "import hnswlib\n",
        "from typing import List, Dict\n",
        "from unstructured.partition.html import partition_html\n",
        "from unstructured.chunking.title import chunk_by_title\n",
        "\n",
        "co = cohere.Client(\"COHERE_API_KEY\") # Get your API key here: https://dashboard.cohere.com/api-keys\n",
        "pc = Pinecone(api_key=\"PINECONE_API_KEY\") # (get API key at app.pinecone.io)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "id": "f0b1935e",
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "import cohere\n",
        "import os\n",
        "co = cohere.Client(os.getenv(\"COHERE_API_KEY\"))\n",
        "pc = Pinecone(\n",
        "    api_key=os.environ.get(\"PINECONE_API_KEY\")\n",
        ")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "id": "Dx1cncziCWBB",
      "metadata": {
        "cellView": "form",
        "id": "Dx1cncziCWBB"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "#@title Enable text wrapping in Google Colab\n",
        "\n",
        "from IPython.display import HTML, display\n",
        "\n",
        "def set_css():\n",
        "  display(HTML('''\n",
        "  <style>\n",
        "    pre {\n",
        "        white-space: pre-wrap;\n",
        "    }\n",
        "  </style>\n",
        "  '''))\n",
        "get_ipython().events.register('pre_run_cell', set_css)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "2f7e7d1c",
      "metadata": {
        "id": "2f7e7d1c"
      },
      "source": [
        "# Create a vector store for ingestion and retrieval\n",
        "\n",
        "First, we define the list of documents we want to ingest and make available for retrieval. As an example, we'll use the contents from the first module of Cohere's *LLM University: What are Large Language Models?*."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "id": "3dca4a88",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        },
        "id": "3dca4a88",
        "outputId": "b05da1ee-0456-4387-c232-a43e0ffed54c"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "raw_documents = [\n",
        "    {\n",
        "        \"title\": \"Text Embeddings\",\n",
        "        \"url\": \"https://docs.cohere.com/docs/text-embeddings\"},\n",
        "    {\n",
        "        \"title\": \"Similarity Between Words and Sentences\",\n",
        "        \"url\": \"https://docs.cohere.com/docs/similarity-between-words-and-sentences\"},\n",
        "    {\n",
        "        \"title\": \"The Attention Mechanism\",\n",
        "        \"url\": \"https://docs.cohere.com/docs/the-attention-mechanism\"},\n",
        "    {\n",
        "        \"title\": \"Transformer Models\",\n",
        "        \"url\": \"https://docs.cohere.com/docs/transformer-models\"}\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5e2a8968",
      "metadata": {
        "id": "5e2a8968"
      },
      "source": [
        "Usually the number of documents for practical applications is vast, and so we'll need to be able to search documents efficiently.  This involves breaking the documents into chunks, generating embeddings, and indexing the embeddings, as shown in the image below.  \n",
        "\n",
        "We implement this in the `Vectorstore` class below, which takes the `raw_documents` list as input.  Three methods are immediately called when creating an object of the `Vectorstore` class:\n",
        "\n",
        "\n",
        "`load_and_chunk()`  \n",
        "This method uses the `partition_html()` method from the `unstructured` library to load the documents from URL and break them into smaller chunks.  Each chunk is turned into a dictionary object with three fields:\n",
        "- `title` - the web page’s title,\n",
        "- `text` - the textual content of the chunk, and\n",
        "- `url` - the web page’s URL.  \n",
        "  \n",
        "  \n",
        "`embed()`  \n",
        "This method uses Cohere's `embed-english-v3.0` model to generate embeddings of the chunked documents.  Since our documents will be used for retrieval, we set `input_type=\"search_document\"`.  We send the documents to the Embed endpoint in batches, because the endpoint has a limit of 96 documents per call.\n",
        "\n",
        "`index()`  \n",
        "This method uses the `hsnwlib` package to index the document chunk embeddings.  This will ensure efficient similarity search during retrieval.  Note that `hnswlib` uses a vector library, and we have chosen it for its simplicity."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 113,
      "id": "7c33412c",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        },
        "id": "7c33412c",
        "outputId": "cf04f8ed-8000-4433-f976-2d37747f21e7"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "class Vectorstore:\n",
        "    \"\"\"\n",
        "    A class representing a collection of documents indexed into a vectorstore.\n",
        "\n",
        "    Parameters:\n",
        "    raw_documents (list): A list of dictionaries representing the sources of the raw documents. Each dictionary should have 'title' and 'url' keys.\n",
        "\n",
        "    Attributes:\n",
        "    raw_documents (list): A list of dictionaries representing the raw documents.\n",
        "    docs (list): A list of dictionaries representing the chunked documents, with 'title', 'text', and 'url' keys.\n",
        "    docs_embs (list): A list of the associated embeddings for the document chunks.\n",
        "    docs_len (int): The number of document chunks in the collection.\n",
        "    idx (hnswlib.Index): The index used for document retrieval.\n",
        "\n",
        "    Methods:\n",
        "    load_and_chunk(): Loads the data from the sources and partitions the HTML content into chunks.\n",
        "    embed(): Embeds the document chunks using the Cohere API.\n",
        "    index(): Indexes the document chunks for efficient retrieval.\n",
        "    retrieve(): Retrieves document chunks based on the given query.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, raw_documents: List[Dict[str, str]]):\n",
        "        self.raw_documents = raw_documents\n",
        "        self.docs = []\n",
        "        self.docs_embs = []\n",
        "        self.retrieve_top_k = 10\n",
        "        self.rerank_top_k = 3\n",
        "        self.load_and_chunk()\n",
        "        self.embed()\n",
        "        self.index()\n",
        "\n",
        "\n",
        "    def load_and_chunk(self) -> None:\n",
        "        \"\"\"\n",
        "        Loads the text from the sources and chunks the HTML content.\n",
        "        \"\"\"\n",
        "        print(\"Loading documents...\")\n",
        "\n",
        "        for raw_document in self.raw_documents:\n",
        "            elements = partition_html(url=raw_document[\"url\"])\n",
        "            chunks = chunk_by_title(elements)\n",
        "            for chunk in chunks:\n",
        "                self.docs.append(\n",
        "                    {\n",
        "                        \"title\": raw_document[\"title\"],\n",
        "                        \"text\": str(chunk),\n",
        "                        \"url\": raw_document[\"url\"],\n",
        "                    }\n",
        "                )\n",
        "\n",
        "    def embed(self) -> None:\n",
        "        \"\"\"\n",
        "        Embeds the document chunks using the Cohere API.\n",
        "        \"\"\"\n",
        "        print(\"Embedding document chunks...\")\n",
        "\n",
        "        batch_size = 90\n",
        "        self.docs_len = len(self.docs)\n",
        "        for i in range(0, self.docs_len, batch_size):\n",
        "            batch = self.docs[i : min(i + batch_size, self.docs_len)]\n",
        "            texts = [item[\"text\"] for item in batch]\n",
        "            docs_embs_batch = co.embed(\n",
        "                texts=texts, model=\"embed-english-v3.0\", input_type=\"search_document\"\n",
        "            ).embeddings\n",
        "            self.docs_embs.extend(docs_embs_batch)\n",
        "\n",
        "    def index(self) -> None:\n",
        "        \"\"\"\n",
        "        Indexes the documents for efficient retrieval.\n",
        "        \"\"\"\n",
        "        print(\"Indexing documents...\")\n",
        "\n",
        "        index_name = 'rag-01'\n",
        "\n",
        "        # If the index does not exist, we create it\n",
        "        if index_name not in pc.list_indexes().names():\n",
        "            pc.create_index(\n",
        "                name=index_name,\n",
        "                dimension=len(self.docs_embs[0]),\n",
        "                metric=\"cosine\",\n",
        "                spec=PodSpec(\n",
        "                    environment=\"gcp-starter\"\n",
        "                )\n",
        "                )\n",
        "    \n",
        "        # connect to index\n",
        "        self.idx = pc.Index(index_name)\n",
        "    \n",
        "        batch_size = 128\n",
        "\n",
        "        ids = [str(i) for i in range(len(self.docs))]\n",
        "        # create list of metadata dictionaries\n",
        "        meta = self.docs\n",
        "\n",
        "        # create list of (id, vector, metadata) tuples to be upserted\n",
        "        to_upsert = list(zip(ids, self.docs_embs, meta))\n",
        "\n",
        "        for i in range(0, len(self.docs), batch_size):\n",
        "            i_end = min(i+batch_size, len(self.docs))\n",
        "            self.idx.upsert(vectors=to_upsert[i:i_end])\n",
        "\n",
        "        # let's view the index statistics\n",
        "        print(\"Indexing complete\")\n",
        "\n",
        "\n",
        "    def retrieve(self, query: str) -> List[Dict[str, str]]:\n",
        "        \"\"\"\n",
        "        Retrieves document chunks based on the given query.\n",
        "\n",
        "        Parameters:\n",
        "        query (str): The query to retrieve document chunks for.\n",
        "\n",
        "        Returns:\n",
        "        List[Dict[str, str]]: A list of dictionaries representing the retrieved document chunks, with 'title', 'text', and 'url' keys.\n",
        "        \"\"\"\n",
        "\n",
        "        docs_retrieved = []\n",
        "        query_emb = co.embed(\n",
        "            texts=[query], model=\"embed-english-v3.0\", input_type=\"search_query\"\n",
        "        ).embeddings\n",
        "\n",
        "\n",
        "        res = self.idx.query(vector=query_emb, top_k=self.retrieve_top_k, include_metadata=True)\n",
        "        docs_to_rerank = [match['metadata']['text'] for match in res['matches']]\n",
        "\n",
        "        rerank_results = co.rerank(\n",
        "            query=query,\n",
        "            documents=docs_to_rerank,\n",
        "            top_n=self.rerank_top_k,\n",
        "            model=\"rerank-english-v2.0\",\n",
        "        )\n",
        "        \n",
        "        docs_reranked = [res['matches'][result.index] for result in rerank_results.results]\n",
        "        \n",
        "        for doc in docs_reranked:\n",
        "            docs_retrieved.append(doc['metadata'])\n",
        "\n",
        "        return docs_retrieved"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e1bf5d85",
      "metadata": {
        "id": "e1bf5d85"
      },
      "source": [
        "In the code cell below, we initialize an instance of the `Vectorstore` class and pass in the `raw_documents` list as input."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 114,
      "id": "4643e630",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 178
        },
        "id": "4643e630",
        "outputId": "fe01fcb6-3574-4322-d8d0-57d37aad397d"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Loading documents...\n",
            "Embedding document chunks...\n",
            "Indexing documents...\n",
            "Indexing complete\n"
          ]
        }
      ],
      "source": [
        "# Create an instance of the Vectorstore class with the given sources\n",
        "vectorstore = Vectorstore(raw_documents)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "61928287",
      "metadata": {
        "id": "61928287"
      },
      "source": [
        "The `Vectorstore` class also has a `retrieve()` method, which we'll use to retrieve relevant document chunks given a query (as in Step 3 in the diagram shared at the beginning of this notebook).  This method has two components: (1) dense retrieval, and (2) reranking.\n",
        "\n",
        "### Dense retrieval\n",
        "\n",
        "First, we embed the query using the same `embed-english-v3.0` model we used to embed the document chunks, but this time we set `input_type=\"search_query\"`.\n",
        "\n",
        "Search is performed by the `knn_query()` method from the `hnswlib` library. Given a query, it returns the document chunks most similar to the query. We can define the number of document chunks to return using the attribute `self.retrieve_top_k=10`.\n",
        "\n",
        "### Reranking\n",
        "\n",
        "After semantic search, we implement a reranking step.  While our semantic search component is already highly capable of retrieving relevant sources, the [Rerank endpoint](https://cohere.com/rerank) provides an additional boost to the quality of the search results, especially for complex and domain-specific queries. It takes the search results and sorts them according to their relevance to the query.\n",
        "\n",
        "We call the Rerank endpoint with the `co.rerank()` method and define the number of top reranked document chunks to retrieve using the attribute `self.rerank_top_k=3`.  The model we use is `rerank-english-v2.0`.  \n",
        "\n",
        "This method returns the top retrieved document chunks `chunks_retrieved` so that they can be passed to the chatbot.\n",
        "\n",
        "In the code cell below, we check the document chunks that are retrieved for the query `\"multi-head attention definition\"`."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "OwozNf_uPEyX",
      "metadata": {
        "id": "OwozNf_uPEyX"
      },
      "source": [
        "## Test Retrieval"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 115,
      "id": "82617b91",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 267
        },
        "id": "82617b91",
        "outputId": "7f1f2bc8-8ed9-4190-bd6b-7af2d9dc1980"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/plain": [
              "[{'text': 'The attention step used in transformer models is actually much more powerful, and it’s called multi-head attention. In multi-head attention, several different embeddings are used to modify the vectors and add context to them. Multi-head attention has helped language models reach much higher levels of efficacy when processing and generating text.',\n",
              "  'title': 'Transformer Models',\n",
              "  'url': 'https://docs.cohere.com/docs/transformer-models'},\n",
              " {'text': \"What you learned in this chapter is simple self-attention. However, we can do much better than that. There is a method called multi-head attention, in which one doesn't only consider one embedding, but several different ones. These are all obtained from the original by transforming it in different ways. Multi-head attention has been very successful at the task of adding context to text. If you'd like to learn more about the self and multi-head attention, you can check out the following two\",\n",
              "  'title': 'The Attention Mechanism',\n",
              "  'url': 'https://docs.cohere.com/docs/the-attention-mechanism'},\n",
              " {'text': 'Attention helps give context to each word, based on the other words in the sentence (or text).',\n",
              "  'title': 'Transformer Models',\n",
              "  'url': 'https://docs.cohere.com/docs/transformer-models'}]"
            ]
          },
          "execution_count": 115,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "vectorstore.retrieve(\"multi-head attention definition\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e69fbca9",
      "metadata": {
        "id": "e69fbca9"
      },
      "source": [
        "# Create a chatbot\n",
        "\n",
        "Next, we implement a class to handle the interaction between the user and the chatbot.  It takes an instance of the `Vectorstore` class as input.\n",
        "\n",
        "The `run()` method will be used to run the chatbot application.  It begins with the logic for getting the user message, along with a way for the user to end the conversation.  \n",
        "\n",
        "Based on the user message, the chatbot needs to decide if it needs to consult external information before responding.  If so, the chatbot determines an optimal set of search queries to use for retrieval.  When we call `co.chat()` with `search_queries_only=True`, the Chat endpoint handles this for us automatically.\n",
        "\n",
        "The generated queries can be accessed from the `search_queries` field of the object that is returned.  Then, what happens next depends on how many queries are returned.\n",
        "- If queries are returned, we call the `retrieve()` method of the Vectorstore object for the  retrieval step.  The retrieved document chunks are then passed to the Chat endpoint by adding a `documents` parameter when we call `co.chat()` again.\n",
        "- Otherwise, if no queries are returned, we call the Chat endpoint another time, passing the user message and without needing to add any documents to the call.\n",
        "\n",
        "In either case, we also pass the `conversation_id` parameter, which retains the interactions between the user and the chatbot in the same conversation thread. We also enable the `stream` parameter so we can stream the chatbot response.\n",
        "\n",
        "We then print the chatbot's response.  In the case that the external information was used to generate a response, we also display citations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 116,
      "id": "d2c15a1f",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        },
        "id": "d2c15a1f",
        "outputId": "8daa9159-338c-45ec-e9ed-830aedcdf0d8"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "class Chatbot:\n",
        "    def __init__(self, vectorstore: Vectorstore):\n",
        "        \"\"\"\n",
        "        Initializes an instance of the Chatbot class.\n",
        "\n",
        "        Parameters:\n",
        "        vectorstore (Vectorstore): An instance of the Vectorstore class.\n",
        "\n",
        "        \"\"\"\n",
        "        self.vectorstore = vectorstore\n",
        "        self.conversation_id = str(uuid.uuid4())\n",
        "\n",
        "    def run(self):\n",
        "        \"\"\"\n",
        "        Runs the chatbot application.\n",
        "\n",
        "        \"\"\"\n",
        "        while True:\n",
        "            # Get the user message\n",
        "            message = input(\"User: \")\n",
        "\n",
        "            # Typing \"quit\" ends the conversation\n",
        "            if message.lower() == \"quit\":\n",
        "              print(\"Ending chat.\")\n",
        "              break\n",
        "            # else:                       # Uncomment for Google Colab to avoid printing the same thing twice\n",
        "              # print(f\"User: {message}\") # Uncomment for Google Colab to avoid printing the same thing twice\n",
        "\n",
        "            # Generate search queries (if any)\n",
        "            response = co.chat(message=message,\n",
        "                               model=\"command-r\",\n",
        "                               search_queries_only=True)\n",
        "\n",
        "            # If there are search queries, retrieve document chunks and respond\n",
        "            if response.search_queries:\n",
        "                print(\"Retrieving information...\", end=\"\")\n",
        "\n",
        "                # Retrieve document chunks for each query\n",
        "                documents = []\n",
        "                for query in response.search_queries:\n",
        "                    documents.extend(self.vectorstore.retrieve(query.text))\n",
        "\n",
        "                # Use document chunks to respond\n",
        "                response = co.chat_stream(\n",
        "                    message=message,\n",
        "                    model=\"command-r\",\n",
        "                    documents=documents,\n",
        "                    conversation_id=self.conversation_id,\n",
        "                )\n",
        "\n",
        "            # If there is no search query, directly respond\n",
        "            else:\n",
        "                response = co.chat_stream(\n",
        "                    message=message,\n",
        "                    model=\"command-r\",\n",
        "                    conversation_id=self.conversation_id,\n",
        "                )\n",
        "\n",
        "            # Print the chatbot response, citations, and documents\n",
        "            print(\"\\nChatbot:\")\n",
        "            citations = []\n",
        "            cited_documents = []\n",
        "\n",
        "            # Display response\n",
        "            for event in response:\n",
        "                if event.event_type == \"text-generation\":\n",
        "                    print(event.text, end=\"\")\n",
        "                elif event.event_type == \"citation-generation\":\n",
        "                    citations.extend(event.citations)\n",
        "                elif event.event_type == \"search-results\":\n",
        "                    cited_documents = event.documents\n",
        "\n",
        "            # Display citations and source documents\n",
        "            if citations:\n",
        "              print(\"\\n\\nCITATIONS:\")\n",
        "              for citation in citations:\n",
        "                print(citation)\n",
        "\n",
        "              print(\"\\nDOCUMENTS:\")\n",
        "              for document in cited_documents:\n",
        "                print(document)\n",
        "\n",
        "            print(f\"\\n{'-'*100}\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "F0X9FZqwOwQ6",
      "metadata": {
        "id": "F0X9FZqwOwQ6"
      },
      "source": [
        "# Run the chatbot"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3cb442f7",
      "metadata": {
        "id": "3cb442f7"
      },
      "source": [
        "We can now run the chatbot.  For this, we create the instance of `Chatbot` and run the chatbot by invoking the `run()` method.\n",
        "\n",
        "The format of each citation is:\n",
        "- `start`: The starting point of a span where one or more documents are referenced\n",
        "- `end`: The ending point of a span where one or more documents are referenced\n",
        "- `text`: The text representing this span\n",
        "- `document_ids`: The IDs of the documents being referenced (`doc_0` being the ID of the first document passed to the `documents` creating parameter in the endpoint call, and so on)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 118,
      "id": "42d3f345",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "42d3f345",
        "outputId": "8b935c8b-b1d4-4913-bdf8-73ba503402b8"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Chatbot:\n",
            "Hello! What's your question? I'm here to help you in any way I can.\n",
            "----------------------------------------------------------------------------------------------------\n",
            "\n",
            "Retrieving information...\n",
            "Chatbot:\n",
            "Word embeddings associate words with lists of numbers, so that similar words are close to each other and dissimilar words are further away.\n",
            "Sentence embeddings do the same thing, but for sentences. Each sentence is associated with a vector of numbers in a coherent way, so that similar sentences are assigned similar vectors, and different sentences are given different vectors.\n",
            "\n",
            "CITATIONS:\n",
            "start=0 end=15 text='Word embeddings' document_ids=['doc_0']\n",
            "start=16 end=53 text='associate words with lists of numbers' document_ids=['doc_0']\n",
            "start=63 end=100 text='similar words are close to each other' document_ids=['doc_0']\n",
            "start=105 end=139 text='dissimilar words are further away.' document_ids=['doc_0']\n",
            "start=140 end=159 text='Sentence embeddings' document_ids=['doc_0', 'doc_2']\n",
            "start=160 end=177 text='do the same thing' document_ids=['doc_0', 'doc_2']\n",
            "start=198 end=211 text='Each sentence' document_ids=['doc_0', 'doc_2']\n",
            "start=215 end=250 text='associated with a vector of numbers' document_ids=['doc_0', 'doc_2']\n",
            "start=256 end=264 text='coherent' document_ids=['doc_2']\n",
            "start=278 end=295 text='similar sentences' document_ids=['doc_0', 'doc_2']\n",
            "start=300 end=324 text='assigned similar vectors' document_ids=['doc_0', 'doc_2']\n",
            "start=330 end=349 text='different sentences' document_ids=['doc_0', 'doc_2']\n",
            "start=354 end=378 text='given different vectors.' document_ids=['doc_0', 'doc_2']\n",
            "\n",
            "DOCUMENTS:\n",
            "{'id': 'doc_0', 'text': 'In the previous chapters, you learned about word and sentence embeddings and similarity between words and sentences. In short, a word embedding is a way to associate words with lists of numbers (vectors) in such a way that similar words are associated with numbers that are close by, and dissimilar words with numbers that are far away from each other. A sentence embedding does the same thing, but associating a vector to every sentence. Similarity is a way to measure how similar two words (or', 'title': 'The Attention Mechanism', 'url': 'https://docs.cohere.com/docs/the-attention-mechanism'}\n",
            "{'id': 'doc_1', 'text': 'Sentence embeddings\\n\\nSo word embeddings seem to be pretty useful, but in reality, human language is much more complicated than simply a bunch of words put together. Human language has structure, sentences, etc. How would one be able to represent, for instance, a sentence? Well, here’s an idea. How about the sums of scores of all the words? For example, say we have a word embedding that assigns the following scores to these words:\\n\\nNo: [1,0,0,0]\\n\\nI: [0,2,0,0]\\n\\nAm: [-1,0,1,0]\\n\\nGood: [0,0,1,3]', 'title': 'Text Embeddings', 'url': 'https://docs.cohere.com/docs/text-embeddings'}\n",
            "{'id': 'doc_2', 'text': 'This is where sentence embeddings come into play. A sentence embedding is just like a word embedding, except it associates every sentence with a vector full of numbers, in a coherent way. By coherent, I mean that it satisfies similar properties as a word embedding. For instance, similar sentences are assigned to similar vectors, different sentences are assigned to different vectors, and most importantly, each of the coordinates of the vector identifies some (whether clear or obscure) property of', 'title': 'Text Embeddings', 'url': 'https://docs.cohere.com/docs/text-embeddings'}\n",
            "\n",
            "----------------------------------------------------------------------------------------------------\n",
            "\n",
            "Retrieving information...\n",
            "Chatbot:\n",
            "The similarities between words and sentences are both quantitative measures of how close the two given items are. There are two types of similarities that can be defined: dot product similarity, and cosine similarity. These methods can determine how similar two words, or sentences, are.\n",
            "\n",
            "CITATIONS:\n",
            "start=54 end=75 text='quantitative measures' document_ids=['doc_0']\n",
            "start=79 end=88 text='how close' document_ids=['doc_0']\n",
            "start=124 end=133 text='two types' document_ids=['doc_0', 'doc_4']\n",
            "start=171 end=193 text='dot product similarity' document_ids=['doc_0', 'doc_4']\n",
            "start=199 end=217 text='cosine similarity.' document_ids=['doc_0', 'doc_4']\n",
            "start=236 end=257 text='determine how similar' document_ids=['doc_0', 'doc_4']\n",
            "\n",
            "DOCUMENTS:\n",
            "{'id': 'doc_0', 'text': 'Now that we know embeddings quite well, let’s move on to using them to find similarities. There are two types of similarities we’ll define in this post: dot product similarity and cosine similarity. Both are very similar and very useful to determine if two words (or sentences) are similar.', 'title': 'Similarity Between Words and Sentences', 'url': 'https://docs.cohere.com/docs/similarity-between-words-and-sentences'}\n",
            "{'id': 'doc_1', 'text': 'But let me add some numbers to this reasoning to make it more clear. Imagine that we calculate similarities for the words in each sentence, and we get the following:\\n\\nThis similarity makes sense in the following ways:\\n\\nThe similarity between each word and itself is 1.\\n\\nThe similarity between any irrelevant word (“the”, “of”, etc.) and any other word is 0.\\n\\nThe similarity between “bank” and “river” is 0.11.\\n\\nThe similarity between “bank” and “money” is 0.25.', 'title': 'The Attention Mechanism', 'url': 'https://docs.cohere.com/docs/the-attention-mechanism'}\n",
            "{'id': 'doc_2', 'text': 'And the results are:\\n\\nThe similarity between sentences 1 and 2: 6738.2858668486715\\n\\nThe similarity between sentences 1 and 3: -122.22666955510499\\n\\nThe similarity between sentences 2 and 3: -3.494608113647928\\n\\nThese results certainly confirm our predictions. The similarity between sentences 1 and 2 is 6738, which is high. The similarities between sentences 1 and 3, and 2 and 3, are -122 and -3.5 (dot products are allowed to be negative too!), which are much lower.', 'title': 'Similarity Between Words and Sentences', 'url': 'https://docs.cohere.com/docs/similarity-between-words-and-sentences'}\n",
            "{'id': 'doc_3', 'text': 'But let me add some numbers to this reasoning to make it more clear. Imagine that we calculate similarities for the words in each sentence, and we get the following:\\n\\nThis similarity makes sense in the following ways:\\n\\nThe similarity between each word and itself is 1.\\n\\nThe similarity between any irrelevant word (“the”, “of”, etc.) and any other word is 0.\\n\\nThe similarity between “bank” and “river” is 0.11.\\n\\nThe similarity between “bank” and “money” is 0.25.', 'title': 'The Attention Mechanism', 'url': 'https://docs.cohere.com/docs/the-attention-mechanism'}\n",
            "{'id': 'doc_4', 'text': 'Now that we know embeddings quite well, let’s move on to using them to find similarities. There are two types of similarities we’ll define in this post: dot product similarity and cosine similarity. Both are very similar and very useful to determine if two words (or sentences) are similar.', 'title': 'Similarity Between Words and Sentences', 'url': 'https://docs.cohere.com/docs/similarity-between-words-and-sentences'}\n",
            "{'id': 'doc_5', 'text': 'And the results are:\\n\\nThe similarity between sentences 1 and 2: 6738.2858668486715\\n\\nThe similarity between sentences 1 and 3: -122.22666955510499\\n\\nThe similarity between sentences 2 and 3: -3.494608113647928\\n\\nThese results certainly confirm our predictions. The similarity between sentences 1 and 2 is 6738, which is high. The similarities between sentences 1 and 3, and 2 and 3, are -122 and -3.5 (dot products are allowed to be negative too!), which are much lower.', 'title': 'Similarity Between Words and Sentences', 'url': 'https://docs.cohere.com/docs/similarity-between-words-and-sentences'}\n",
            "\n",
            "----------------------------------------------------------------------------------------------------\n",
            "\n",
            "Ending chat.\n"
          ]
        }
      ],
      "source": [
        "# Create an instance of the Chatbot class\n",
        "chatbot = Chatbot(vectorstore)\n",
        "\n",
        "# Run the chatbot\n",
        "chatbot.run()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "-JBWZVz9ObcV",
      "metadata": {
        "id": "-JBWZVz9ObcV"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.4"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}
